# Lists to store training and validation losses
train_loss = []
val_loss = []
# Lists to store SSIM and PSNR values
train_secret_ssim = []
train_cover_ssim = []
train_secret_psnr = []
train_cover_psnr = []
val_secret_ssim = []
val_cover_ssim = []
val_secret_psnr = []
val_cover_psnr = []
def calculate_psnr(true, pred, max_val=1.0):
return tf.image.psnr(true, pred, max_val=max_val)
for epoch in range(EPOCHS):
epoch_train_loss = []
epoch_train_secret_ssim = []
epoch_train_cover_ssim = []
epoch_train_secret_psnr = []
epoch_train_cover_psnr = []
print(f"Epoch {epoch+1}/{EPOCHS}")
for step in range(len(train_files) // BATCH_SIZE):
cover_batch, secret_batch = next(train_data)
if np.any(np.isnan(cover_batch)) or np.any(np.isnan(secret_batch)):
print(f'NaN values found in data batch at epoch {epoch+1}, step {step+1}.')
continue
with tf.GradientTape() as tape:
# Forward pass
hiding_output, reveal_output = model(secret_batch, cover_batch)
# Calculate loss
loss, secret_ssim, cover_ssim, cover_dist = get_loss_op(
secret_batch, reveal_output, cover_batch, hiding_output, beta=BETA, gamma=GAMMA)
if tf.reduce_any(tf.math.is_nan(loss)):
print(f'NaN loss encountered at epoch {epoch+1}, step {step+1}.')
continue
# Backpropagation
gradients = tape.gradient(loss, model.trainable_variables)
clipped_gradients, _ = tf.clip_by_global_norm(gradients, 1.0) # Adjust clip_norm as needed
optimizer.apply_gradients(zip(clipped_gradients, model.trainable_variables))
epoch_train_loss.append(loss.numpy())
epoch_train_secret_ssim.append(secret_ssim.numpy())
epoch_train_cover_ssim.append(cover_ssim.numpy())
epoch_train_secret_psnr.append(calculate_psnr(secret_batch, reveal_output).numpy())
epoch_train_cover_psnr.append(calculate_psnr(cover_batch, hiding_output).numpy())
train_loss.append(np.mean(epoch_train_loss))
train_secret_ssim.append(np.mean(epoch_train_secret_ssim))
train_cover_ssim.append(np.mean(epoch_train_cover_ssim))
train_secret_psnr.append(np.mean(epoch_train_secret_psnr))
train_cover_psnr.append(np.mean(epoch_train_cover_psnr))
# Validation
epoch_val_loss = []
epoch_val_secret_ssim = []
epoch_val_cover_ssim = []
epoch_val_secret_psnr = []
epoch_val_cover_psnr = []
for step in range(len(val_files) // BATCH_SIZE):
cover_batch, secret_batch = next(val_data)
if np.any(np.isnan(cover_batch)) or np.any(np.isnan(secret_batch)):
print(f'NaN values found in validation data batch at epoch {epoch+1}, step {step+1}.')
continue
# Forward pass (no gradient calculation during validation)
hiding_output, reveal_output = model(secret_batch, cover_batch)
loss, secret_ssim, cover_ssim, cover_dist = get_loss_op(secret_batch, reveal_output, cover_batch, hiding_output, beta=BETA, gamma=GAMMA)
if tf.reduce_any(tf.math.is_nan(loss)):
print(f'NaN loss encountered in validation at epoch {epoch+1}, step {step+1}.')
continue
epoch_val_loss.append(loss.numpy())
epoch_val_secret_ssim.append(secret_ssim.numpy())
epoch_val_cover_ssim.append(cover_ssim.numpy())
epoch_val_secret_psnr.append(calculate_psnr(secret_batch, reveal_output).numpy())
epoch_val_cover_psnr.append(calculate_psnr(cover_batch, hiding_output).numpy())
val_loss.append(np.mean(epoch_val_loss))
val_secret_ssim.append(np.mean(epoch_val_secret_ssim))
val_cover_ssim.append(np.mean(epoch_val_cover_ssim))
val_secret_psnr.append(np.mean(epoch_val_secret_psnr))
val_cover_psnr.append(np.mean(epoch_val_cover_psnr))
print(f'Epoch {epoch+1}/{EPOCHS}, Train Loss: {train_loss[-1]}, Val Loss: {val_loss[-1]}')
print(f'Epoch {epoch+1}/{EPOCHS}, Train SSIM (Secret): {train_secret_ssim[-1]}, Val SSIM (Secret): {val_secret_ssim[-1]}')
print(f'Epoch {epoch+1}/{EPOCHS}, Train SSIM (Cover): {train_cover_ssim[-1]}, Val SSIM (Cover): {val_cover_ssim[-1]}')
print(f'Epoch {epoch+1}/{EPOCHS}, Train PSNR (Secret): {train_secret_psnr[-1]}, Val PSNR (Secret): {val_secret_psnr[-1]}')
print(f'Epoch {epoch+1}/{EPOCHS}, Train PSNR (Cover): {train_cover_psnr[-1]}, Val PSNR (Cover): {val_cover_psnr[-1]}')
# Test the model on test data and show the loss
test_loss = []
for step in range(len(test_files) // BATCH_SIZE):
cover_batch, secret_batch = next(test_data)
hiding_output, reveal_output = model(secret_batch, cover_batch)
loss, secret_ssim, cover_ssim, cover_dist = get_loss_op(secret_batch, reveal_output, cover_batch, hiding_output, beta=BETA, gamma=GAMMA)
test_loss.append(loss.numpy())
print(f'Test Loss: {np.mean(test_loss)}')
